#include "videoPush.h"

Video::Video()
    : isRunning(false),
      input_ctx(nullptr),
      input_codec_ctx(nullptr),
      swsContext(nullptr),
      output_ctx(nullptr),
      output_codec_ctx(nullptr)
{
}

void Video::init()
{
    init_input();
    init_output();
    swsContext = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_YUYV422, WIDTH, HEIGHT, AV_PIX_FMT_YUV420P, SWS_BILINEAR, nullptr, nullptr, nullptr);
}

void Video::start()
{
    isRunning = true;

    recvThread = std::thread(&Video::recv_work_fun, this);
    convertThread = std::thread(&Video::convert_work_fun, this);
    sendThread = std::thread(&Video::send_work_fun, this);
}

bool Video::getIsRunning() const
{
    return isRunning.load();
}

Video::~Video()
{
    isRunning = false;
    recvThread.join();
    convertThread.join();
    sendThread.join();

    avcodec_free_context(&input_codec_ctx);
    avcodec_free_context(&output_codec_ctx);

    avformat_close_input(&input_ctx);
    avformat_free_context(input_ctx);
    avformat_free_context(output_ctx);

    sws_freeContext(swsContext);
}

void Video::recv_work_fun()
{
    AVPacket *pkt = av_packet_alloc();
    int ret;
    while (isRunning && av_read_frame(input_ctx, pkt) >= 0)
    {
        ret = avcodec_send_packet(input_codec_ctx, pkt);
    }
    av_packet_free(&pkt);
    isRunning.store(false);
}

void Video::convert_work_fun()
{
    AVFrame *input_frame = av_frame_alloc();
    // 申请Frame
    AVFrame *output_frame = av_frame_alloc();
    output_frame->format = AV_PIX_FMT_YUV420P;
    output_frame->width = WIDTH;
    output_frame->height = HEIGHT;
    output_frame->pts = input_frame->pts;
    // 根据帧信息申请Frame.data
    av_frame_get_buffer(output_frame, 0);
    int ret;
    while (isRunning)
    {
        ret = avcodec_receive_frame(input_codec_ctx, input_frame);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF || ret < 0)
        {
            std::cout << "error 1" << std::endl;
            continue;
        }
        sws_scale(swsContext, input_frame->data, input_frame->linesize, 0, input_frame->height, output_frame->data, output_frame->linesize);

        output_frame->pts = input_frame->pts;
        ret = avcodec_send_frame(output_codec_ctx, output_frame);
    }
    av_frame_free(&input_frame);
    av_frame_free(&output_frame);
}

void Video::send_work_fun()
{
    AVPacket *pkt = av_packet_alloc();
    int ret;
    while (isRunning)
    {
        ret = avcodec_receive_packet(output_codec_ctx, pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF || ret < 0)
        {
            std::cout << "error 2" << std::endl;
            continue;
        }
        av_interleaved_write_frame(output_ctx, pkt);
        // av_packet_unref(pkt);
    }
    av_packet_free(&pkt);
}

void Video::init_input()
{
    avdevice_register_all();
    auto input_fmt = av_find_input_format("video4linux2");
    avformat_open_input(&input_ctx, "/dev/video4", input_fmt, nullptr);
    avformat_find_stream_info(input_ctx, nullptr);
    int video_stream_index = -1;
    for (int i = 0; i < input_ctx->nb_streams; i++)
    {
        if (input_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            video_stream_index = i;
            break;
        }
    }
    input_codec_ctx = avcodec_alloc_context3(nullptr);
    avcodec_parameters_to_context(input_codec_ctx, input_ctx->streams[video_stream_index]->codecpar);
    auto input_codec = avcodec_find_decoder(input_codec_ctx->codec_id);
    avcodec_open2(input_codec_ctx, input_codec, nullptr);
}

void Video::init_output()
{
    avformat_network_init();
    // 获取输出流
    output_ctx = nullptr;
    avformat_alloc_output_context2(&output_ctx, nullptr, "rtsp", "rtsp://192.168.0.60:8554/stream");
    auto output_stream = avformat_new_stream(output_ctx, nullptr);
    // 获取输出编解码器
    auto output_codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    output_codec_ctx = avcodec_alloc_context3(output_codec);
    output_codec_ctx->width = WIDTH;
    output_codec_ctx->height = HEIGHT;
    output_codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
    output_codec_ctx->time_base = {1, FPS};
    output_codec_ctx->bit_rate = 1000000;

    // 赋值输出编码器参数到输出流
    avcodec_parameters_from_context(output_stream->codecpar, output_codec_ctx);
    avcodec_open2(output_codec_ctx, output_codec, nullptr);
    // 写入输出头文件
    avformat_write_header(output_ctx, nullptr);
}
